{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TR-nF8awMM1A",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Using [Jina](https://github.com/jina-ai/jina) on Colab with GPU/TPU support\n",
    "\n",
    "In this tutorial, we will cover two ways of using Colab as illustrated below:\n",
    "\n",
    "![](https://raw.githubusercontent.com/jina-ai/jina/master/docs/how-to/jina-on-colab.svg)\n",
    "\n",
    "## 1. Change runtime type\n",
    "\n",
    "Go to menu `Runtime -> Change run time type -> GPU/TPU`\n",
    "\n",
    "\n",
    "## 2. Install Jina\n",
    "You will be asked to restart the kernel after installation. Go ahead and click \"Restart *Runtime*\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yzG5IfrxKt-v",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "!pip install jina"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BoRWbFUHN-4y",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 3. Import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KaKzywlEKv3q",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from jina import Flow, Document, Executor, requests, DocumentArray\n",
    "import jina\n",
    "\n",
    "print(jina.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hlwpiW_hOPBt",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 4. Build a GPU Executor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "p58WUjjgMgFy",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "class GPUExec(Executor):\n",
    "\n",
    "  @requests\n",
    "  def foo(self, docs: DocumentArray, **kwargs):\n",
    "    docs[0].tags['cuda'] = torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "u2wjIBQXOXrv",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 5. Add GPU Executor to a Flow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3q35cMikLnyA",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "f = Flow().add(uses=GPUExec)\n",
    "\n",
    "f.start()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_GROx0BnOgdN",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 6. Send request to the running Flow\n",
    "\n",
    "The result should contain a single Document which has `.tags['cuda']` field set to `True` via `GPUExec`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ACoTXhn-Lpz5",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "r = f.post('/', Document())\n",
    "print(r[0].tags)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0K3bFdStR09C",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now we see that one can directly connect to a running Flow inside Google Colab. The trick here is to use `f.start()` to start the Flow and use `f.close()` to close the Flow. This usage is different than our classic `with` context manager."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FJ6l4bsgLyqz",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "f.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hBDjAUkoSNJW",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Connecting from local to a running Flow on Colab"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sEm8mxdKSa0k",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "A more interesting use case is to connect to Colab from a **local** machine. In this case, we need to install `ngrox`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OvebGiRzXmc4",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "!pip install pyngrok"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0ASjGLBhXono",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 1. Expose via HTTP\n",
    "\n",
    "\n",
    "Easiest and most compatible, no registration needed on ngrox. However, HTTP protocol is not performant on large data and streaming."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UeXfAwBuSsdY",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "f = Flow(protocol='http', port=54321).add(uses=GPUExec)\n",
    "\n",
    "f.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5DTM1cQOVSOI",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "!ngrok http 54321 --log \"stdout\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OsKJXEETYIRN",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Notice the last line, you should see something like the following:\n",
    "\n",
    "```\n",
    "t=2022-06-11T20:06:49+0000 lvl=info msg=\"started tunnel\" obj=tunnels name=command_line addr=http://localhost:54321 url=https://7ee9-35-197-36-214.ngrok.io\n",
    "```\n",
    "\n",
    "Now use the string after `url=` to replace the following `curl` command. In this example, my `url=https://7ee9-35-197-36-214.ngrok.io`. **Yours would be different!**\n",
    "\n",
    "```bash\n",
    "curl -X POST https://7ee9-35-197-36-214.ngrok.io/post \\\n",
    "     -H 'Content-Type: application/json' \\\n",
    "     -d '{\"data\":[{}]}'\n",
    "```\n",
    "\n",
    "Running this curl command in your local terminal. You will get:\n",
    "\n",
    "```json\n",
    "{\"header\":{\"requestId\":\"008c15ac9770412b9ea8fcd3b0944520\",\"status\":null,\"execEndpoint\":\"/\",\"targetExecutor\":\"\"},\"parameters\":null,\"routes\":[{\"executor\":\"gateway\",\"startTime\":\"2022-06-11T20:09:30.189579+00:00\",\"endTime\":\"2022-06-11T20:09:30.227964+00:00\",\"status\":null},{\"executor\":\"executor0\",\"startTime\":\"2022-06-11T20:09:30.190086+00:00\",\"endTime\":\"2022-06-11T20:09:30.227796+00:00\",\"status\":null}],\"data\":[{\"id\":\"402bddfd49311dab4128dcfae5cec9ba\",\"parent_id\":null,\"granularity\":null,\"adjacency\":null,\"blob\":null,\"tensor\":null,\"mime_type\":null,\"text\":null,\"weight\":null,\"uri\":null,\"tags\":{\"cuda\":true},\"offset\":null,\"location\":null,\"embedding\":null,\"modality\":null,\"evaluations\":null,\"scores\":null,\"chunks\":null,\"matches\":null}]}\n",
    "```\n",
    "\n",
    "Where one can see `\"tags\":{\"cuda\":true}` is successfully returned.\n",
    "\n",
    "To stop the server, click stop the last cell. Then close the Flow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iTc0d_Urb9nW",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "f.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wom9p0hJb5f7",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 2. Expose via gRPC\n",
    "\n",
    "Using ngrox to forward gRPC is also possible. But you will need to first sign up at https://dashboard.ngrok.com/signup\n",
    "\n",
    "After signing up, you can get a token. Then simply add your token via (replacing `YOUR_TOKEN_HERE`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mGsPDHFIcTL7",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "!ngrok authtoken 2ARf0Y7QrG3E2TsFaX7W3KWvfGD_6R972KKvbqeucpCrCuEHv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "x7nX4DA3cy2q",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now let's start the Flow with gRPC gateway again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eZWf3cJwUDkM",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "f = Flow(port=54321).add(uses=GPUExec)\n",
    "\n",
    "f.start()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KzrbTL3cdIxb",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Let's get its public address via `ngrok`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "y-dV3792c7K1",
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "!ngrok tcp 54321 --log \"stdout\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lYxxAvHtdPY8",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "At the last line, you should see something like: \n",
    "\n",
    "```\n",
    "t=2022-06-11T20:29:11+0000 lvl=info msg=\"started tunnel\" obj=tunnels name=command_line addr=//localhost:54321 url=tcp://6.tcp.ngrok.io:18096\n",
    "```\n",
    "\n",
    "Grab the text after `url=tcp://` in my case it is `6.tcp.ngrok.io:18096`.\n",
    "\n",
    "Now build a client using this address from your local laptop/Python environment.\n",
    "\n",
    "Copy paste the code below to your local Python, remmeber to change your address.\n",
    "\n",
    "```python\n",
    "from jina import Client, Document\n",
    "\n",
    "\n",
    "c = Client(host='grpc://6.tcp.ngrok.io:18096')\n",
    "r = c.post('/', Document())\n",
    "print(r[0].tags['cuda'])\n",
    "```\n",
    "\n",
    "And you will get \n",
    "\n",
    "```text\n",
    "True\n",
    "```\n",
    "\n",
    "Showing the connection is success!\n",
    "\n",
    "Now enjoy the free GPU/TPU to build your awesome Jina applications!"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Using Jina on Colab.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}